/*
 * Copyright (c) 2015 Minshenglife. All rights reserved. Minshenglife
 * PROPRIETARY/CONFIDENTIAL. Use is subject to license terms.
 */
package com.sentry.commons.crypto;

import com.sentry.commons.utils.Exceptions;

import org.apache.commons.lang3.ArrayUtils;

import java.security.GeneralSecurityException;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;

/**
 * AES工具类，链式调用，方式如下：<br>
 * <pre>
 * AES aes = AES.create().cbc().pkcs7Padding().iv(iv).key(key);
 * byte[] enc = aes.encrypt(text.getBytes());
 * </pre>
 * 或者使用默认快速方式
 * <pre>
 * AES aes = AES.create(true).key(key);
 * </pre>
 * @author WuBo (wub1@minshenglife.com)
 *
 */
public class AES {



    private static final String AES = "AES";

    private static final String ECB = "ECB";
    private static final String CBC = "CBC";

    private static final String NOPADDING = "NOPADDING";
    private static final String PKCS5PADDING = "PKCS5PADDING";
    private static final String PKCS7PADDING = "PKCS7PADDING";

    private AES() {}

    /**
     * 默认AES方式，AES/ECB/PKCS5PADDING
     * @param _default
     * @return
     */
    public static AES create(boolean _default) {
        AES aes = new AES();
        if(_default){
            aes.mode = ECB;
            aes.padding = PKCS5PADDING;
        }
        return aes;
    }
    
    public static AES create() {
        return create(false);
    }    


    private String mode;
    private String padding;
    private byte[] key;
    private byte[] iv;
    
    
    
    public AES cbc() {
        this.mode = CBC;
        return this;
    }

    public AES ecb() {
        this.mode = ECB;
        return this;
    }

    public AES noPadding() {
        this.padding = NOPADDING;
        return this;
    }

    public AES pkcs5Padding() {
        this.padding = PKCS5PADDING;
        return this;
    }

    public AES pkcs7Padding() {
        this.padding = PKCS7PADDING;
        return this;
    }

    public AES key(byte[] key) {
        this.key = key;
        return this;
    }

    public AES iv(byte[] iv) {
        this.iv = iv;
        return this;
    }

    public byte[] encrypt(byte[] input) {
        if (CBC.equals(this.mode) && this.iv != null) {
            return aes(input, this.key, this.iv, Cipher.ENCRYPT_MODE);
        } else {
            return aes(input, this.key, Cipher.ENCRYPT_MODE);
        }
    }

    public byte[] decrypt(byte[] input) {
        if (CBC.equals(this.mode) && this.iv != null) {
            return aes(input, this.key, this.iv, Cipher.DECRYPT_MODE);
        } else {
            return aes(input, this.key, Cipher.DECRYPT_MODE);
        }
    }



    private byte[] aes(byte[] input, byte[] key, int mode) {
        return aes(input, key, null, mode);
    }



    /**
     * 使用AES加密或解密无编码的原始字节数组, 返回无编码的字节数组结果.
     * 
     * @param input 原始字节数组
     * @param key 符合AES要求的密钥
     * @param iv 初始向量
     * @param mode Cipher.ENCRYPT_MODE 或 Cipher.DECRYPT_MODE
     */
    private byte[] aes(byte[] input, byte[] key, byte[] iv, int mode) {
        try {
            SecretKey secretKey = new SecretKeySpec(key, AES);
            Cipher cipher = null;
            byte[] inputData = input;
            String padding = this.padding;;
            if (PKCS7PADDING.equals(this.padding)) {
                if(Cipher.ENCRYPT_MODE == mode){
                    inputData = padding(input);
                }
                padding = NOPADDING;
            }


            String transName = "AES/" + this.mode + "/" + padding;
            if (iv != null) {
                IvParameterSpec ivSpec = new IvParameterSpec(iv);
                cipher = Cipher.getInstance(transName);
                cipher.init(mode, secretKey, ivSpec);
            } else {
                cipher = Cipher.getInstance(transName);
                cipher.init(mode, secretKey);
            }
            
            byte[] result = cipher.doFinal(inputData);
            if (PKCS7PADDING.equals(this.padding) && Cipher.DECRYPT_MODE == mode) {
                result = dePadding(cipher.doFinal(inputData));
            }
            
            return result;
        } catch (GeneralSecurityException e) {
            throw Exceptions.unchecked(e);
        }
    }

    // pkcs7padding
    private byte[] padding(final byte[] input) {
        byte[] pa = PKCS7PaddingEncoder.encode(input.length);
        return ArrayUtils.addAll(input, pa);
    }
    
    private byte[] dePadding(final byte[] input) {
        return PKCS7PaddingEncoder.decode(input);
    }


}
